package top.wangjingxin.util;

import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.transport.client.PreBuiltTransportClient;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.stereotype.Component;
import top.wangjingxin.model.article.ArticleDTO;
import top.wangjingxin.model.article.SearchDataVO;
import top.wangjingxin.model.article.SearchVO;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static java.lang.Integer.parseInt;
import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;


/**
 * Created by 王镜鑫 on 17-11-23.
 * 请关注一下他的个人博客 wangjingxin.top
 */
@Component
public class ElasticUtil implements DisposableBean {

    private TransportClient client;
    private HighlightBuilder hiBuilder = new HighlightBuilder();

    public ElasticUtil() {
        try {
            client = new PreBuiltTransportClient(Settings.EMPTY)
                    .addTransportAddress(new TransportAddress(InetAddress.getByName("localhost"), 9300));
            hiBuilder.preTags("<span style='color:red'>");
            hiBuilder.postTags("</span>");
            hiBuilder.field("c");
        } catch (UnknownHostException e) {
            e.printStackTrace();
        }
    }

    public SearchVO query(String key, Page page) {
        SearchVO result = new SearchVO();
        List<SearchDataVO> data = new ArrayList<>();
        result.setData(data);
        SearchResponse response = client.prepareSearch("aipage")
                .setTypes("c")
                .setSearchType(SearchType.QUERY_THEN_FETCH)
                .setQuery(QueryBuilders.matchQuery("c", key))
                .setFrom(page.getPage()).setSize(page.getRows())
                .setFetchSource(false)
                .highlighter(hiBuilder)
                .get();
        result.setCount(response.getHits().getTotalHits());
        response.getHits().iterator().forEachRemaining(hit->{
            SearchDataVO vo = new SearchDataVO();
            vo.setId(parseInt(hit.getId()));
            vo.setContent(t2s(hit.getHighlightFields().get("c").getFragments()));
            data.add(vo);
        });
        return result;
    }

    private String t2s(Text[] cs) {
        StringBuilder sb = new StringBuilder();
        Arrays.stream(cs).forEach(c-> sb.append("...").append(c.string()));
        return sb.append("...").toString();
    }

    public int insert(ArticleDTO dto) {
        try {
            IndexResponse response = client.prepareIndex("aipage", "c", dto.getId() + "")
                    .setSource(jsonBuilder()
                            .startObject()
                            .field("c", dto.getPlainContent())
                            .endObject())
                    .get();
            return response.status().getStatus();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return 1;
    }

    public long total() {
        return client.prepareSearch("aipage")
                .setTypes("c")
                .get().getHits().getTotalHits();
    }

    @Override
    public void destroy() {
        client.close();
    }

    public static void main(String[] args) {
        ElasticUtil util = new ElasticUtil();
        System.out.println(util.query("奔驰宝马",new Page(1,10)));
    }
}